import torch
import cv2 as cv
from model.yolov5 import YoloV5
from model.decode import decode_box, scale_coords, plot_one_box
from utils.utils import letter_box_ractangle
import time


model = YoloV5()
model.load_state_dict(torch.load('yolov5.pth'))
image_origin = cv.imread('E:/bus.jpg')
assert image_origin is not None, 'Image does not exit.'
image = letter_box_ractangle(image_origin)
image = image[:, :, ::-1].transpose(2, 0, 1)/255
print('Network loading complete.')
model.eval()
with torch.no_grad():
    t1 = time.time()
    image = torch.tensor(image, dtype=torch.float32).unsqueeze(0)
    predict = model(image)
    predict = decode_box(predict)
    print(time.time() - t1)
det = predict[0]
det[:, :4] = scale_coords(image.shape[2:], det[:, :4], image_origin.shape).round()
for *xyxy, conf, cls in reversed(det):
    label = f'{int(cls)} {conf:.4f}'
    plot_one_box(xyxy, image_origin, label=label, color=(0, 0, 255), line_thickness=1)
cv.imshow('test', image_origin)
cv.waitKey(0)
# cv.imwrite('test.png', image_origin)
